#!/usr/bin/python
version = '1.63'
"""converts Wuerzburg data to PCST data, calls dhea and does everything 
version = '1.63'

Author:     Gunnar W. Klau
Version:    see global version variable


History:
            29 Oct 09, 1.63 fixed bug when doing suboptimal solutions
	     4 Aug 09, 1.62 recompiled dhea, some fixes in Bool options
	     4 Aug 09, 1.61 output also in matrix format (one file only)
            31 Jul 09, 1.6  re-added multiple weights support (bootstrapping) 
            17 Mar 09, 1.51 added DIMACS graph format export (for further treewidth analysis by A. Koster)
             9 Dec 08, 1.5  added permutations again (nodes *and* edges)
             5 Dec 08, 1.4  options for kantenheinz, knotenheinz, kantenknotenheinz
             3 Dec 08, 1.3  only *one* heinz script. removed permutations for the time being. asym 0 is again default (in the latest test it seemed again to be better)
            18 Nov 08, 1.2  supports edge weights. asym 1 is again default (seemed to work better)  
             1 Oct 08, 1.18 new dhea version (asym 0 works now). this is now
			    the new default since it seems to improve running
			    times significantly (not tested on really large 
			    instances, though)
	    22 Sep 08, 1.17 added default options, improved option handling 
             7 Dec 07, 1.16 changed options to heuristics on again. works much better
                            for our instances. 
            30 Nov 07, 1.15 added option -a and changed default settings for dhea
                            now, heuristics are turned off (see e-mail Ivana)
            21 Sep 07, 1.10 better option support (guwek)
    	    18 Sep 07, 1.00 created (guwek)
"""

### schoener machen:
## * environment-var ILM_ACCESS_FILE ueberpruefen
## * check whether sol dir exists

## * min_node_weight not initialized to float ("infinity"). problem?

import sys
import re
import os
import glob
from optparse import OptionParser
import traceback

n = 0 # number of nodces
min_node_weight = {} # the smallest negative weight 
sum_node_weights = {} # sum of all transformed weights
m = 0 # number of edges
nodes = {}
edges = {}

node_weights = {}
orig_node_weights = {}
orig_edge_weights = {}
nodeidx = {}
verbose = False

edge_weights = {}
sum_edge_weights = {}
min_edge_weight = {}

min_weight = {}

no_scores = 1 


def parse_nodes(nodeFile):
    """Parse the nodes from the file containing the nodes
    """
    # global variables modified by this function:
    global n 
    global min_node_weight 
    global node_weights
    global orig_node_weights
    global no_scores
    
    try:
       i = 1	      
       line = nodeFile.readline() # read header line
       while (line[0] == "#"): line = nodeFile.readline() # eat comments
       while line:
	 L = line.split();
	 no_scores = len(L)-1
	 nodes[i] = L[0]
     
         for j in range (1, no_scores + 1):
	   if consider_node_weights == True:
             orig_node_weights[(i, j)] = node_weights[(i, j)] = float(L[j].replace(",", "."))
           else:
             orig_node_weights[(i, j)] = node_weights[(i, j)] = 0
           if not min_node_weight.has_key(j): min_node_weight[j] = float("infinity")
           min_node_weight[j] = min(min_node_weight[j], node_weights[(i, j)])
	 nodeidx[L[0]] = i
         i = i + 1
         line = nodeFile.readline()
       n = i - 1
    except:
       print 'Failed parsing nodes'
       traceback.print_exc()
       return

def parse_edges(edgeFile):
    """Parse the edges from the file containing the edges
    """
    global m 
    global min_edge_weight
    global edge_weights
    global orig_edge_weights

    try:
       i = 1	
       line = edgeFile.readline() # read header line
       while (line[0] == "#"): line = edgeFile.readline() # eat comments
       while line:
	 L = line.split()
         if L[0] == L[1]: # we remove self loops here
           if verbose: print "<WARNING> skipping self loop (", L[0], ",", L[1], ")"
         else: 
           if not nodeidx.has_key(L[0]) or not nodeidx.has_key(L[1]): 
              if verbose and not nodeidx.has_key(L[0]): print "<WARNING> node ", L[0], " not in nodefile" 
              if verbose and not nodeidx.has_key(L[1]): print "<WARNING> node ", L[1], " not in nodefile" 
           else:
    	     edges[i] = [nodeidx[L[0]], nodeidx[L[1]]]
             for j in range (1, no_scores + 1):
               if consider_edge_weights == True:
                 print "setting ", i, j, " to ", float(L[j+1].replace(",", "."))
                 if len(L) >= j+2: orig_edge_weights[(i, j)] = float(L[j+1].replace(",", "."))   
                 else: orig_edge_weights[(i, j)] = orig_edge_weights[(i, j-1)]
               else:
                 orig_edge_weights[(i, j)] = 0
               if not generate_instances:
                   edge_weights[(i, j)] = orig_edge_weights[(i, j)] # positive because our edge weights **are** costs
               else:
                   edge_weights[(i, j)] = -orig_edge_weights[(i, j)] # negative because our edge weights are *not* costs 
                
               if not min_edge_weight.has_key(j): min_edge_weight[j] = float("infinity")
	       min_edge_weight[j] = min(min_edge_weight[j], edge_weights[(i, j)])
             i = i + 1
         line = edgeFile.readline()
       m = i - 1
    except:
       print 'Failed parsing edges'
       traceback.print_exc()
       return

def export_DIMACS(nodefilestr, edgefilestr):
    global n
    global node_weights
    global nodeidx
    global min_node_weight
    global sum_node_weights
    global m
    global edge_weights
    global min_edge_weight
    global sum_edge_weights
    global min_weight
    global nodes
    global edges

    """Read the input files."""
    try:
        print "exporting to DIMACS format..."
        #Open and parse input files.
        nodeFile = open(nodefilestr)
        edgeFile = open(edgefilestr)
        
        parse_nodes(nodeFile)
        parse_edges(edgeFile)
        nodeFile.close()
	edgeFile.close()

        filename = nodefilestr + ".dimacs"
        outputFile = open(filename, 'w')
        #Write some comment
        outputFile.write('c protein-protein interaction graph')
        outputFile.write('\nc created by heinz version ' + version)
        outputFile.write('\nc author: Gunnar W. Klau (Centrum Wiskunde & Informatica)')
        outputFile.write('\nc -------------------------------------------------------')
        outputFile.write('\np mnwcs ' + str(n) + ' '  + str(m))
        #for i in range(1,n+1):
         # outputFile.write(str(i) + "\t0\t0\t" + str(node_weights[(i, j)]) + "\n")
        #outputFile.write('link\n')
        for i in range(1,m+1):
          outputFile.write('\ne\t' + str(edges[i][0]) + "\t" + str(edges[i][1]))
        outputFile.close()


    except:
        print 'problem parsing input'
        traceback.print_exc()
        #Put here some more information - usage...




def wue2stein(nodefilestr, edgefilestr, heinztmpstr):
    global n
    global node_weights
    global nodeidx
    global min_node_weight
    global sum_node_weights
    global m
    global edge_weights
    global min_edge_weight
    global sum_edge_weights
    global min_weight
    global nodes
    global edges

    """Read the input files."""
    try:
        #Open and parse input files.
        nodeFile = open(nodefilestr)
        edgeFile = open(edgefilestr)
        
        parse_nodes(nodeFile)
        parse_edges(edgeFile)
        nodeFile.close()
	edgeFile.close()
    except:
        print 'problem parsing input'
        traceback.print_exc()
        #Put here some more information - usage...

    #transform:
    # now transform the weights and compute their sum
    for j in range (1, no_scores + 1):
      sum_node_weights[j] = sum_edge_weights[j] = 0
      min_weight[j] = min(min_edge_weight[j], min_node_weight[j])
      for i in range(1, n+1):
        node_weights[(i, j)] -= min_weight[j]
        sum_node_weights[j] += node_weights[(i, j)]
      for i in range(1, m+1):
        edge_weights[(i, j)] -= min_weight[j]
        sum_edge_weights[j] += edge_weights[(i, j)]

    try:
      for j in range (1, no_scores+1):
        filename = heinztmpstr + "." + str(j) + ".pst"
        outputFile = open(filename, 'w')
        #Write some comment
        outputFile.write('#created by heinz (' + str(n) + ' nodes, ' + str(m) + ' edges)\n')  
        outputFile.write('#sum of all node weights: ' + str(sum_node_weights[j]))
        outputFile.write('\n#sum of all edge weights: ' + str(sum_edge_weights[j]))
        outputFile.write('\n#min node weight: ' + str(min_node_weight[j]))
        outputFile.write('\n#min edge weight: ' + str(min_edge_weight[j]))
        outputFile.write('\n#min weight: ' + str(min_weight[j]))
        outputFile.write('\nnode\n')
        for i in range(1,n+1):
          outputFile.write(str(i) + "\t0\t0\t" + str(node_weights[(i, j)]) + "\n")
        outputFile.write('link\n')
        for i in range(1,m+1):
          outputFile.write(str(i) + "\t" + str(edges[i][0]) + "\t" + str(edges[i][1]) + "\t" + str(edge_weights[(i, j)]) + "\n")
        outputFile.close()

    except:
	print 'problem writing output'
        traceback.print_exc()



def directStein(nodefilestr, edgefilestr, heinztmpstr):
    global n
    global node_weights
    global nodeidx
    global min_node_weight
    global sum_node_weights
    global m
    global edge_weights
    global min_edge_weight
    global sum_edge_weights
    global min_weight
    global nodes
    global edges

    """Read the input files."""
    try:
        #Open and parse input files.
        nodeFile = open(nodefilestr)
        edgeFile = open(edgefilestr)
        
        parse_nodes(nodeFile)
        parse_edges(edgeFile)
        nodeFile.close()
        edgeFile.close()
    except:
        print 'problem parsing input'
        traceback.print_exc()
        #Put here some more information - usage...

    #transform:
    # now transform the weights and compute their sum
    for j in range (1, no_scores + 1):
        sum_node_weights[j] = sum_edge_weights[j] = 0
        min_weight[j] = min(min_edge_weight[j], min_node_weight[j])
        for i in range(1, n+1):    
            #node_weights[(i, j)] -= min_weight[j]
            sum_node_weights[j] += node_weights[(i, j)]
        for i in range(1, m+1): 
            #edge_weights[(i, j)] -= min_weight[j]
            sum_edge_weights[j] += edge_weights[(i, j)]

    try:
      for j in range (1, no_scores+1):
        filename = heinztmpstr + "." + str(j) + ".pst"
        outputFile = open(filename, 'w')
        #Write some comment
        outputFile.write('#created by heinz (' + str(n) + ' nodes, ' + str(m) + ' edges)\n')  
        outputFile.write('#sum of all node weights: ' + str(sum_node_weights[j]))
        outputFile.write('\n#sum of all edge weights: ' + str(sum_edge_weights[j]))
        outputFile.write('\n#min node weight: ' + str(min_node_weight[j]))
        outputFile.write('\n#min edge weight: ' + str(min_edge_weight[j]))
        outputFile.write('\n#min weight: ' + str(min_weight[j]))
        outputFile.write('\nnode\n')
        for i in range(1,n+1):
          outputFile.write(str(i) + "\t0\t0\t" + str(node_weights[(i, j)]) + "\n")
        outputFile.write('link\n')
        for i in range(1,m+1):
          outputFile.write(str(i) + "\t" + str(edges[i][0]) + "\t" + str(edges[i][1]) + "\t" + str(edge_weights[(i, j)]) + "\n")
        outputFile.close()

    except:
        print 'problem writing output'
        traceback.print_exc()





#And now: The program.
parser = OptionParser()
parser.add_option("-n", "--nodefile", dest="nodeFile", help="WUE file with node information", metavar="FILE", default="few_nodes.txt")
parser.add_option("-e", "--edgefile", dest="edgeFile", help="WUE file with edge information", metavar="FILE", default="few_edges.txt")
parser.add_option("-s", "--no_subopt", dest="no_subopt", help="number of suboptimal solutions", default=1)
parser.add_option("-d", "--subopt_diff", dest="subopt_diff", help="difference between sol (Ham. dist in % or -1)", default=-1)
parser.add_option("-t", "--tolerance", dest="tolerance", help="solutions may be this percentage off optimum (default: 0)", default=0)
parser.add_option("-H", "--heinztmp", dest="heinztmp", help="heinz tmp dir", default="heinztmp")
parser.add_option("-g", "--generate", dest="generate_instances", help="do first step (generate the transformed instances)", default="True")
parser.add_option("-r", "--run", dest="run_heinz", help="run PCST code on transformed instances", default="True")
parser.add_option("-c", "--collect", dest="collect_solutions", help="collect solutions and write heinz outputfile", default="True")
parser.add_option("-a", "--additional", dest="additional_dhea_options", help="additional options for dhea", default="")
parser.add_option("-v", "--verbose", dest="verbose", help="verbose output", action="store_true")
parser.add_option("-N", "--consider_node_weights", dest="consider_node_weights", help="consider node weights", default="True")
parser.add_option("-E", "--consider_edge_weights", dest="consider_edge_weights", help="consider edge weights", default="True")
parser.add_option("-D", "--DIMACS", dest="DIMACS", help="convert to DIMACS format", default=False, action="store_true")

(options, args) = parser.parse_args()

#print options

nodeFile = options.nodeFile
edgeFile = options.edgeFile
no_subopt = options.no_subopt
subopt_diff = options.subopt_diff
tolerance = options.tolerance
heinztmp = options.heinztmp
generate_instances = eval(options.generate_instances)
#generate_instances = True
run_heinz = eval(options.run_heinz)
collect_solutions = eval(options.collect_solutions)
additional_dhea_options = options.additional_dhea_options
verbose = options.verbose
consider_node_weights = eval(options.consider_node_weights)
consider_edge_weights = eval(options.consider_edge_weights)
DIMACS = options.DIMACS

if DIMACS == True:
    consider_edge_weights = False
    export_DIMACS(str(nodeFile), str(edgeFile))
    sys.exit()



if generate_instances == True:  
    if os.path.exists(heinztmp):   
        print "deleting old tmp dir... ",
        os.popen("rm -rf " + heinztmp)
        print "done."
    os.popen("mkdir " + heinztmp)
    
    print "transforming subnet problem instance into prize-collecting Steiner tree... ",
    sys.stdout.flush() 
    
    wue2stein(str(nodeFile), str(edgeFile), str(heinztmp) + "/heinz.tmp")
    print "done."
    
else:
    if os.path.exists(heinztmp):   
        print "deleting old tmp dir... ",
        os.popen("rm -rf " + heinztmp)
        print "done."
    os.popen("mkdir " + heinztmp)
    
    print "Heinz is going to use the original node and edge files"
    sys.stdout.flush()
    directStein(str(nodeFile), str(edgeFile), str(heinztmp) + "/heinz.tmp")
    print "done."
      
if run_heinz == True:
  if verbose: verb_string = ""
  else: verb_string = "2>/dev/null"
  for heinz in glob.glob(heinztmp + "/heinz*"):
    cmd_str = "./dhea ifile " + heinz + " solutiondir " + heinztmp + " odir " + heinztmp + " tolerance " + str(tolerance) + \
    " no_subopt " + str(no_subopt) + " subopt_diff " + str(subopt_diff) + " asym 0 " + additional_dhea_options + verb_string

    print "running " + cmd_str + "... ",
    sys.stdout.flush() 
    os.popen(cmd_str)
    print "done."



if collect_solutions == True: 
  print "transforming prize-collecting Steiner tree solution into subnetwork solution... ",
  
  solfilelist = glob.glob(heinztmp + "/*.solution")
  no_solutions = len(solfilelist)
  #print "no of solutions: ", no_solutions
  #print "no of scores * no of subopt: ", int(no_scores) * int(no_subopt)
  
  if int(no_solutions) != int(no_scores) * int(no_subopt): 
    print "<ERROR> number of solutions does not equal number of scores times number of suboptimal solutions."
    exit(-1)

  dheasolfile = {}
  #for j in range(1, no_solutions + 1): dheasolfile[i] = open(heinztmp + "/"

 
  #for i in range(5, len(sys.argv)): solFile[i-4]  = open(sys.argv[i], 'r')
  
  for k in range (0, int(no_subopt)):
      noderesultfilename = nodeFile + "." + str(k) + ".hnz"
      edgeresultfilename = edgeFile + "." + str(k) + ".hnz"

      outputfile_node = open(noderesultfilename, 'w')
      outputfile_edge = open(edgeresultfilename, 'w')
      
      outputfile_node.write("#label\t")
      for j in range (1, no_scores + 1): outputfile_node.write('\tscore' + str(j))
      outputfile_node.write("\n")
      
      
      outputfile_edge.write("#nodeA  \tnodeB")
      for j in range (1, no_scores + 1): outputfile_edge.write('\tscore' + str(j))
      outputfile_edge.write("\n")
      
      node_in_solution = {} 
      edge_in_solution = {} 
      z_node = {}
      z_edge = {}

      for j in range(1, no_scores + 1):

          dheasolfilename = heinztmp + "/heinz.tmp." + str(j) + ".pst." + str(k) + ".solution"
          dheasolfile = open(dheasolfilename, 'r')
          
          z_node[j] = z_edge[j] = 0
          for i in range(1, n+1): node_in_solution[(i,j)] = False
          for i in range(1, m+1): edge_in_solution[(i,j)] = False

          line = dheasolfile.readline()
          while line[0:3] != '===': # we are in the node section
              L = line.split()
              node_in_solution[(eval(L[0]), j)] = True
              z_node[j] += orig_node_weights[(eval(L[0]), j)]
              line = dheasolfile.readline()
          line = dheasolfile.readline()
          while line[0:3] != '===': # we are in the link section
              L = line.split()
              edge_in_solution[(eval(L[0]), j)] = True
              z_edge[j] += orig_edge_weights[(eval(L[0]), j)]
              line = dheasolfile.readline()

      for i in range(1, n+1):
          outputfile_node.write(nodes[i])
          outputfile_node.write("\t")
          for j in range(1, no_scores + 1):
              if node_in_solution[(i, j)] == True: outputfile_node.write(str(orig_node_weights[(i, j)]) + "\t")
              else: outputfile_node.write("n/a\t")
          outputfile_node.write("\n")
      outputfile_node.write("#subnet node score")
      outputfile_node.write("\t")
      for j in range (1, no_scores + 1): outputfile_node.write(str(z_node[j]) + "\t")
      outputfile_node.write("\n")
      outputfile_node.close()

      for i in range(1, m+1):
          outputfile_edge.write(nodes[edges[i][0]] + "\t" + nodes[edges[i][1]])
          outputfile_edge.write("\t")
          for j in range(1, no_scores + 1):
              if edge_in_solution[(i,j)] == True: outputfile_edge.write(str(orig_edge_weights[i, j]) + "\t")
              else: outputfile_edge.write("n/a          ")
          outputfile_edge.write("\n")
      outputfile_edge.write("#subnet edge score")
      outputfile_edge.write("\t")
      for j in range (1, no_scores + 1): outputfile_edge.write(str(z_edge[j]) + "\t")
      outputfile_edge.write("\n")
      outputfile_edge.close()
  print "done."

